import threading
from channels.generic.websocket import WebsocketConsumer
from django_redis import get_redis_connection
from asgiref.sync import async_to_sync
from channels.exceptions import StopConsumer
from threading import Thread
from system.host.models import Host
import json


# class ExecConsumer(WebsocketConsumer):
#     def __init__(self, *args, **kwargs):
#         super().__init__(*args, **kwargs)
#         self.token = self.scope['url_route']['kwargs']['token']
#         self.rds = get_redis_connection()
#
#     def connect(self):
#         self.accept()
#
#     def disconnect(self, code):
#         self.rds.close()
#
#     def get_response(self):
#         response = self.rds.brpop(self.token, timeout=5)
#         return response[1] if response else None
#
#     def receive(self, **kwargs):
#         response = self.get_response()
#         while response:
#             data = response.decode()
#             self.send(text_data=data)
#             response = self.get_response()
#         self.send(text_data='pong')


class ExecConsumer(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.token = self.scope['url_route']['kwargs']['token']
        self.rds = get_redis_connection()

    def websocket_connect(self, message):
        print('connect')
        self.accept()
        # 将相同token来的人放入到一个群中
        async_to_sync(self.channel_layer.group_add)(self.token, self.channel_name)

    def websocket_disconnect(self, message):
        self.rds.close()
        async_to_sync(self.channel_layer.group_discard)(self.token, self.channel_name)
        raise StopConsumer()

    def websocket_receive(self, message):
        thread = threading.Thread(target=self.deploy)
        thread.start()

    def get_response(self):
        response = self.rds.brpop(self.token)
        return response[1] if response else None

    def deploy(self):
        response = self.get_response()
        while response:
            data = response.decode()
            async_to_sync(self.channel_layer.group_send)(self.token, {'type': 'my.send', 'message': data})
            response = self.get_response()
        async_to_sync(self.channel_layer.group_send)(self.token, {'type': 'my.send', 'message': 'pong'})

    def my_send(self, event):
        message = event['message']
        self.send(message)


class SSHConsumer(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.user = self.scope['user']
        self.id = self.scope['url_route']['kwargs']['id']
        self.chan = None
        self.ssh = None

    def loop_read(self):
        while True:
            data = self.chan.recv(32 * 1024)
            # print('read: {!r}'.format(data))
            if not data:
                self.close(3333)
                break
            self.send(bytes_data=data)

    def receive(self, text_data=None, bytes_data=None):
        data = text_data or bytes_data
        if data and self.chan:
            data = json.loads(data)
            # print('write: {!r}'.format(data))
            resize = data.get('resize')
            if resize and len(resize) == 2:
                self.chan.resize_pty(*resize)
            else:
                self.chan.send(data['data'])

    def disconnect(self, code):
        self.chan.close()
        self.ssh.close()
        # print('Connection close')

    def connect(self):
        # if has_host_perm(self.user, self.id):
        self.accept()
        self._init()
        # else:
        # self.close()

    def _init(self):
        self.send(bytes_data=b'\r\33[KConnecting ...\r')
        host = Host.objects.filter(pk=self.id).first()
        if not host:
            self.send(text_data='Unknown host\r\n')
            self.close()
        try:
            self.ssh = host.get_ssh().get_client()
        except Exception as e:
            self.send(bytes_data=f'Exception: {e}\r\n'.encode())
            self.close()
            return
        self.chan = self.ssh.invoke_shell(term='xterm')
        self.chan.transport.set_keepalive(30)
        Thread(target=self.loop_read).start()
